use crate::solution::Solution;
// use std::collections::HashMap;
use std::collections::HashSet;

impl Solution {
    pub fn arithmetic_triplets(nums: Vec<i32>, diff: i32) -> i32 {
        let (mut cnt, mut cache) = (0, HashSet::new());
        nums.into_iter().for_each(|num| {
            if cache.contains(&(num - diff)) && cache.contains(&(num - 2 * diff)) { cnt += 1 };
            cache.insert(num);
        });
        cnt

    }

}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_1() {
        // example usage
        let arr = [1, 2, 3, 4, 5].to_vec();
        let diff = 2;
        println!("{}", Solution::arithmetic_triplets(arr, diff)); // output: 4
    }
}